feat: add per-head fp8 kv attention calibration by zeel2104 · Pull Request #2044 · intel/auto-round

zeel2104 · 2026-07-10T15:47:19Z

Resolve #2026

Description

Adds per-head FP8 KV/attention calibration support on top of the existing per-tensor static FP8 paths.

Main changes:

Adds static_kv_granularity and static_attention_granularity options with "tensor" as the default and "head" as the new per-head mode.
Collects per-head k_scale, v_scale, and q_scale for attention tensors shaped as [batch, heads, seq, head_dim].
Exports LLMC metadata with attn_head strategy for per-head KV/attention scales.
Preserves existing per-tensor behavior by default.
Adds tests for helper MHA/GQA scale shapes, LLMC export, and AutoRound export.
Fixes a Windows-only test file-lock issue by closing safe_open handles before reusing the output path.

Type of Change

New feature

Related Issues

Relates to #2026

Checklist Before Submitting

My code has been tested locally.
Documentation has been updated as needed.
New or updated tests are included where applicable.
The CUDA CI has passed. You can trigger it by commenting /azp run Unit-Test-CUDA-AutoRound.

Tested locally:

python -m pytest test/test_cpu/test_experimental_fp8_granularity.py -q
python -m pytest test/test_cpu/export/test_llmc_format.py -q
python -m pytest test/test_cpu/export/test_export.py -k "per_head or static_afp8_export or static_fp8_attn" -q
python -m pytest test/test_cpu/quantization/test_mxfp_nvfp.py -k "fp8_kv_attn" -q

Note:

test/test_cpu/integrations/test_inc_integration.py was not run because neural_compressor is not installed locally.

zeel2104 · 2026-07-10T15:47:54Z

/azp run Unit-Test-CUDA-AutoRound

azure-pipelines · 2026-07-10T15:48:00Z

Commenter does not have sufficient privileges for PR 2044 in repo intel/auto-round

yiliu30 · 2026-07-12T05:56:12Z

/azp run Unit-Test-CUDA-AutoRound

azure-pipelines · 2026-07-12T05:56:23Z

Azure Pipelines successfully started running 1 pipeline(s).

yiliu30 · 2026-07-14T05:35:21Z

/azp run Unit-Test-CUDA-AutoRound

azure-pipelines · 2026-07-14T05:35:31Z

Azure Pipelines successfully started running 1 pipeline(s).

yiliu30 · 2026-07-24T09:17:03Z

Hi @zeel2104, thanks for the contribution! Overall, it looks good to me. Could you please rebase the changes onto the latest main branch and resolve the CI issues?

We also need to quantize the model and validate inference in vLLM using the Triton backend.

yiliu30 · 2026-07-29T01:07:37Z

Hi @zeel2104, are you still working on this PR? We're targeting this feature for the current release. If you're tied up with other work, please let me know. Thanks!

Signed-off-by: Zeel <desaizeel2128@gmail.com>

zeel2104 · 2026-07-29T22:00:01Z

Hi @yiliu30, sorry for late reply. I rebased this PR onto the latest main and pushed the updated branch.

I also resolved the LLMC exporter conflict by keeping the newer top-level attention_input_activations schema and applying the per-head attn_head strategy there.

Local validation after rebase:

python -m pytest test/test_cpu/test_experimental_fp8_granularity.py -q
python -m pytest test/test_cpu/export/test_llmc_format.py -k "kv_config or attention_config" -q --tb=short
python -m pytest test/test_cpu/export/test_export.py -k "per_head or static_afp8_export or static_fp8_attn" -q --tb=short

I do not have vLLM/Triton available locally, so I’ll need CI or a CUDA environment for the vLLM Triton inference validation.

yiliu30

Findings

[P1] auto_round/experimental/utils.py:69: The new per-head FP8 QDQ path casts directly through torch.float8_e4m3fn instead of using the project’s STE helper. Tensor granularity uses float8_e4m3fn_ste() in auto_round/data_type/fp8.py, which preserves gradient flow during optimized calibration. With static_*_granularity="head" and iters > 0, this new path can cut or mishandle gradients through Q/K/V, so per-head mode is not equivalent to the existing static FP8 behavior. Please route the per-head cast through the same STE helper and add a gradient-flow test.

yiliu30 · 2026-07-31T08:57:45Z

Hi @yiliu30, sorry for late reply. I rebased this PR onto the latest main and pushed the updated branch.

I also resolved the LLMC exporter conflict by keeping the newer top-level attention_input_activations schema and applying the per-head attn_head strategy there.

Local validation after rebase:

python -m pytest test/test_cpu/test_experimental_fp8_granularity.py -q

python -m pytest test/test_cpu/export/test_llmc_format.py -k "kv_config or attention_config" -q --tb=short

python -m pytest test/test_cpu/export/test_export.py -k "per_head or static_afp8_export or static_fp8_attn" -q --tb=short

I do not have vLLM/Triton available locally, so I’ll need CI or a CUDA environment for the vLLM Triton inference validation.

Hi @zeel2104 it is fine, could you quantize the Qwen3-8B and upload it the hugging face. I can help verify it.

…attn # Conflicts: # auto_round/autoround.py # auto_round/compressors/entry.py # test/test_cpu/export/test_export.py # test/test_cpu/export/test_llmc_format.py

yiliu30 · 2026-08-05T04:23:57Z

/azp run Unit-Test-CUDA-AutoRound

azure-pipelines · 2026-08-05T04:24:07Z

Azure Pipelines successfully started running 1 pipeline(s).

wenhuach21 · 2026-08-05T06:45:32Z

+        "--static_kv_granularity",
+        default="tensor",
+        type=str,
+        choices=["tensor", "head"],


We already have too many CLI arguments. Is there a way to group all the KV-related options into a kv_scheme and provide some predefined presets to cover the common use cases?

Yes, I’m also thinking about refining the KV-cache and attention quantization configs into dedicated schemes, something like:

"kv_cache_scheme": { "dynamic": false, "group_size": null, "num_bits": 8, "strategy": "tensor" ... }

I’ll prepare an RFC for this later. For this PR, though, I think the current functionality looks good.

wenhuach21 · 2026-08-05T06:49:27Z

        self.is_immediate_saving = is_immediate_saving
        self.static_kv_dtype = static_kv_dtype
        self.static_attention_dtype = static_attention_dtype
+        self.static_kv_granularity = static_kv_granularity


Personally, I think KV-related options belong in the scheme context rather than the compression context.

@n1ck-guo, could you refine the definition of each context and clearly explain its purpose? That would make it much easier for developers to understand which context a new option should belong to.

…attn # Conflicts: # test/unit/test_cpu/test_experimental_fp8_granularity.py

yiliu30 · 2026-08-07T06:07:11Z

/azp run Unit-Test-CUDA-AutoRound

azure-pipelines · 2026-08-07T06:07:21Z

Azure Pipelines successfully started running 1 pipeline(s).

for more information, see https://pre-commit.ci

yiliu30 · 2026-08-07T07:39:55Z

/azp run Unit-Test-CUDA-AutoRound

azure-pipelines · 2026-08-07T07:40:05Z

Azure Pipelines successfully started running 1 pipeline(s).

zeel2104 force-pushed the feat/fp8-per-head-kv-attn branch 2 times, most recently from 43cd2cb to d2d95fb Compare July 10, 2026 17:21

feat: add per-head fp8 kv attention calibration

4b672fc

Signed-off-by: Zeel <desaizeel2128@gmail.com>

zeel2104 force-pushed the feat/fp8-per-head-kv-attn branch from bfc6e01 to 4b672fc Compare July 29, 2026 21:52

yiliu30 requested changes Jul 31, 2026

View reviewed changes

Merge remote-tracking branch 'origin/main' into feat/fp8-per-head-kv-…

ec6cb05

…attn # Conflicts: # auto_round/autoround.py # auto_round/compressors/entry.py # test/test_cpu/export/test_export.py # test/test_cpu/export/test_llmc_format.py

yiliu30 self-requested a review August 5, 2026 04:23

yiliu30 requested review from lvliang-intel, n1ck-guo and wenhuach21 August 5, 2026 05:19

wenhuach21 reviewed Aug 5, 2026

View reviewed changes

Merge branch 'main' into feat/fp8-per-head-kv-attn

ec402e6

yiliu30 force-pushed the feat/fp8-per-head-kv-attn branch from 686cb08 to ec402e6 Compare August 7, 2026 02:28

yiliu30 added 2 commits August 7, 2026 02:29

Merge remote-tracking branch 'origin/main' into feat/fp8-per-head-kv-…

c99f852

…attn # Conflicts: # test/unit/test_cpu/test_experimental_fp8_granularity.py

Merge branch 'main' into feat/fp8-per-head-kv-attn

1fffef6

lvliang-intel reviewed Aug 7, 2026

View reviewed changes

Comment thread test/unit/test_cpu/export/test_export.py

lvliang-intel reviewed Aug 7, 2026

View reviewed changes

Comment thread auto_round/experimental/utils.py

lvliang-intel reviewed Aug 7, 2026

View reviewed changes

Comment thread auto_round/experimental/utils.py

yiliu30 and others added 2 commits August 7, 2026 07:33

fix: address fp8 review comments

e9766a5

[pre-commit.ci] auto fixes from pre-commit.com hooks

3a3aa43

for more information, see https://pre-commit.ci

yiliu30 approved these changes Aug 7, 2026

View reviewed changes

Uh oh!

Conversation

zeel2104 commented Jul 10, 2026 • edited by yiliu30 Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Type of Change

Related Issues

Checklist Before Submitting

Uh oh!

zeel2104 commented Jul 10, 2026

Uh oh!

azure-pipelines Bot commented Jul 10, 2026

Uh oh!

yiliu30 commented Jul 12, 2026

Uh oh!

azure-pipelines Bot commented Jul 12, 2026

Uh oh!

yiliu30 commented Jul 14, 2026

Uh oh!

azure-pipelines Bot commented Jul 14, 2026

Uh oh!

yiliu30 commented Jul 24, 2026

Uh oh!

yiliu30 commented Jul 29, 2026

Uh oh!

zeel2104 commented Jul 29, 2026

Uh oh!

yiliu30 left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

yiliu30 commented Jul 31, 2026

Uh oh!

yiliu30 commented Aug 5, 2026

Uh oh!

azure-pipelines Bot commented Aug 5, 2026

Uh oh!

wenhuach21 Aug 5, 2026

Choose a reason for hiding this comment

Uh oh!

yiliu30 Aug 7, 2026

Choose a reason for hiding this comment

Uh oh!

wenhuach21 Aug 5, 2026

Choose a reason for hiding this comment

Uh oh!

yiliu30 commented Aug 7, 2026

Uh oh!

azure-pipelines Bot commented Aug 7, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

yiliu30 commented Aug 7, 2026

Uh oh!

azure-pipelines Bot commented Aug 7, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

zeel2104 commented Jul 10, 2026 •

edited by yiliu30

Loading

yiliu30 left a comment •

edited

Loading